import torch

# CLIP模型配置
CLIP_CONFIG = {
    'model_name': 'openai/clip-vit-base-patch32',  # 预训练模型名称
    'image_size': 224,                             # 输入图像尺寸
    'text_max_length': 77,                         # 文本最大长度
    'device': 'cuda' if torch.cuda.is_available() else 'cpu',  # 设备选择
}

# 训练配置
TRAIN_CONFIG = {
    'batch_size': 4,              # 批次大小 (减小以适应GPU内存)
    'learning_rate': 1e-4,         # 学习率
    'num_epochs': 50,              # 训练轮数
    'warmup_steps': 1000,          # 预热步数
    'weight_decay': 0.1,           # 权重衰减
    'max_grad_norm': 1.0,          # 梯度裁剪
    'device': 'cuda' if torch.cuda.is_available() else 'cpu',  # 自动选择设备
    'model_name': 'openai/clip-vit-base-patch32',  # 预训练模型名称
}

# 数据配置
DATA_CONFIG = {
    'image_dir': 'data/images',    # 图像数据目录
    'text_file': 'data/captions.txt',  # 文本描述文件
    'train_split': 0.8,            # 训练集比例
    'val_split': 0.1,              # 验证集比例
    'test_split': 0.1,             # 测试集比例
}

# 保存配置
SAVE_CONFIG = {
    'model_save_path': 'models/clip_model.pth',    # 模型保存路径
    'log_save_path': 'logs/training_log.csv',      # 训练日志路径
    'checkpoint_dir': 'checkpoints',               # 检查点目录
}